import os
import fitz
from PIL import Image


# pdf转图片
def pdf_to_images(pdf_path):
    print('开始转换pdf文件：' + pdf_path)
    # 获取输出文件夹名称
    output_path = os.path.splitext(pdf_path)[0]
    # 创建输出文件夹
    os.makedirs(output_path, exist_ok=True)
    # 打开PDF文件
    with fitz.open(pdf_path) as doc:
        # 遍历每个页面
        for page_number in range(len(doc)):
            # 获取页面对象
            page = doc.load_page(page_number)
            # 获取页面的像素图
            pix = page.get_pixmap()
            # 转换PDF页面为Pillow图像
            image = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
            # 构建输出文件路径
            image_path = f"{output_path}/page_{page_number}.png"
            # 将Pillow图像保存为文件
            image.save(image_path)
    print('转换完成：' + output_path)
    return output_path


def pdf_to_images_import(pdf_path):
    print('开始转换pdf文件：' + pdf_path)
    # 获取输出文件夹名称
    output_path = os.path.splitext(pdf_path)[0]
    # 创建输出文件夹
    os.makedirs(output_path, exist_ok=True)
    image_urls = []
    # 打开PDF文件
    with fitz.open(pdf_path) as doc:
        # 遍历每个页面
        for page_number in range(len(doc)):
            # 获取页面对象
            page = doc.load_page(page_number)
            # 获取页面的像素图
            pix = page.get_pixmap()
            # 转换PDF页面为Pillow图像
            image = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
            # 构建输出文件路径
            image_path = f"{output_path}/page_{page_number}.png"
            # 将Pillow图像保存为文件
            image.save(image_path)
            image_urls.append(image_path)
    print('转换完成：' + output_path)
    return image_urls
